[TRTLLM-15264][test] Kimi K3 disagg review fixups: KDA test geometry, gate docs, example cleanup - #17483
[TRTLLM-15264][test] Kimi K3 disagg review fixups: KDA test geometry, gate docs, example cleanup#17483brnguyen2 wants to merge 9 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughKimi K3 KDA disaggregated-serving documentation now describes sharded and replicated layouts. Transfer tests derive production cache parameters and cover matched or heterogeneous TP with attention-DP enabled or disabled. ChangesKDA transfer validation
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/disaggregated/test_kda_mamba_transfer.py`:
- Around line 691-716: Add unittest/disaggregated/test_kda_mamba_transfer.py to
the repository’s integration test registration list so test_kda_transfer and the
other module tests are executed by CI and QA; update the existing test-list
configuration rather than changing the test implementation.
- Around line 527-544: Unset UCX_NET_DEVICES during setup of the KDA loopback
transfer tests, before creating any NIXL transceivers, while preserving the
existing UCX_TLS and TRTLLM_NIXL_NUM_THREADS configuration. Apply this to the
shared setup covering test_kda_hetero_tp_sharded_accepted and test_kda_transfer
so CI-provided device selection cannot affect either test.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1a1dfd1b-a513-412a-b4d7-7402c910563f
📒 Files selected for processing (8)
examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yamlexamples/kimi_k3/disagg/README.mdexamples/kimi_k3/disagg/ctx_config.yamltensorrt_llm/_torch/disaggregation/native/bounce/config.pytensorrt_llm/_torch/disaggregation/native/bounce/impl.pytensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.pytensorrt_llm/_torch/disaggregation/transceiver.pytests/unittest/disaggregated/test_kda_mamba_transfer.py
💤 Files with no reviewable changes (1)
- examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml
|
/bot run |
|
PR_Github #65160 [ run ] triggered by Bot. Commit: |
|
PR_Github #65160 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65196 [ run ] triggered by Bot. Commit: |
|
PR_Github #65196 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65220 [ run ] triggered by Bot. Commit: |
|
PR_Github #65220 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65280 [ run ] triggered by Bot. Commit: |
|
PR_Github #65280 [ run ] completed with state
|
|
/bot run |
1da35af to
f6098c7
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
PR_Github #65322 [ run ] triggered by Bot. Commit: |
|
PR_Github #65322 [ run ] completed with state
|
|
PR_Github #66561 [ run ] triggered by Bot. Commit: |
|
PR_Github #66561 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66563 [ run ] triggered by Bot. Commit: |
|
PR_Github #66563 [ run ] completed with state
|
…arison-overlap) Two pre-existing mypy errors surface on PRs whose build compiles the bindings, triggering the full mypy type check: - _meet_max_token_stop_criteria returned an Any-typed boolean expression (operands come from Any-typed C++ binding attributes) from a -> bool function [no-any-return]; wrap the return in bool(). - update_requests compared req.state to GENERATION_COMPLETE after an earlier 'continue' narrowed that literal away. The state is mutated at runtime by the intervening calls (process_draft_tokens -> _handle_stop_criteria -> finish_by), which mypy cannot track, so the comparison is valid; annotate with a precise type: ignore and comment [comparison-overlap]. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
PR_Github #66564 [ run ] triggered by Bot. Commit: |
|
PR_Github #66564 [ run ] completed with state
|
PR NVIDIA#17609 de-enrolled a batch of perf/test_perf_sanity.py::test_e2e DeepSeek-V3.2 / Kimi-K2.5 cases from the test-db perf_sanity lists but left their entries in waives.txt. Those 18 entries now reference tests that exist in no L0 or QA list, so the "Check Test List" stage fails on every PR with "Non-existent test name in l0 or qa list found in waives.txt". Remove the orphaned entries. Their nvbug SKIPs are moot now that the tests no longer run. Verified: no perf_sanity waive references a non-existent list entry after this change; AST validation and the duplicate-waives check both pass. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com> (cherry picked from commit c42fdcd)
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #66581 [ run ] triggered by Bot. Commit: |
|
PR_Github #66582 [ run ] triggered by Bot. Commit: |
|
PR_Github #66581 [ run ] completed with state |
|
PR_Github #66582 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66633 [ run ] triggered by Bot. Commit: |
|
PR_Github #66633 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66638 [ run ] triggered by Bot. Commit: |
|
PR_Github #66638 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66803 [ ] completed with state |
|
/bot run |
|
PR_Github #66839 [ run ] triggered by Bot. Commit: |
|
PR_Github #66839 [ run ] completed with state |
5 participants
Description
Follow-ups from review of the Kimi K3 disaggregated-serving work, bundled into one change:
_create_kda_managersintest_kda_mamba_transfer.pybuilds its Mamba cache params viaextract_mamba_kv_cache_paramswith unscaled global head counts, letting the manager's own attention-DP gate do the sharding — matching how the model runtime head-shards KDA state across TP when attention-DP is off.validate_peer_compatiblecorrectly accepts hetero-TP peers;test_kda_hetero_tp_rejectedbecametest_kda_hetero_tp_sharded_accepted, asserting concrete per-rank slot sizes before acceptance. New transfer parametrizations cover matched TP2 (attention-DP on and off) and hetero 2->4 / 4->2 with bitwise comparison.peer.pydocstring/error text, the K3 disagg example README, andctx_config.yaml— the claim that hetero ctx/gen TP is rejected for K3 predated the head-sharded layout. Matched geometry remains the only configuration validated end-to-end.bounce/config.pyexplaining what each admission control covers.transceiver.py_slice_num_bytes(sums all pool views) andbounce/impl.pyblock_bytes_per_group(pool 0 only), pointing at the TRTLLM-15194 sizing gap without changing behavior.Test Coverage
On Blackwell hardware:
tests/unittest/disaggregated/test_kda_mamba_transfer.py18 passed;tests/unittest/disaggregated/test_bounce.py75 passed. Product-code changes are comment/docstring/error-text only.PR Checklist
[TRTLLM-15264][test]conventionDev Engineer Review
extract_mamba_kv_cache_params.QA Engineer Review
Modified test code:
tests/unittest/disaggregated/test_kda_mamba_transfer.pytest_kda_hetero_tp_rejected()withtest_kda_hetero_tp_sharded_accepted().test_kda_transfer(ctx_tp, gen_tp)to acceptenable_attention_dp.tests/unittest/disaggregated/test_mamba_transfer.py_create_transceiverswithenable_attention_dp.Mapping.These tests are not individually listed in
tests/integration/test_lists/. The change adds one waiver entry intests/integration/test_lists/waives.txt, but it does not provide CI or manual-QA entries for the modified test functions.Verdict: needs follow-up. Confirm CBTS coverage and add or verify the required CI and manual-QA test-list coverage.